import os
from langchain_community.document_loaders import PyPDFLoader
from langchain_community.document_loaders import UnstructuredMarkdownLoader
from langchain_community.document_loaders import Docx2txtLoader
from langchain_community.document_loaders import TextLoader

class DocumentLoader:
    """文档加载器，支持多种格式文件的加载"""
    def __init__(self):
        # 初始化支持的文件加载器映射
        self.loader_mapping = {
            ".pdf": PyPDFLoader,
            ".docx": Docx2txtLoader,
            ".txt": TextLoader,
            ".md": UnstructuredMarkdownLoader,
        }

    def load_single_document(self, file_path: str) -> list:
        """
        加载单个文档
        :param file_path: 文件路径
        :return: 文档列表
        """
        if not os.path.exists(file_path):
            raise FileNotFoundError(f"文件不存在: {file_path}")

        # 获取文件扩展名
        ext = os.path.splitext(file_path)[1].lower()
        if ext not in self.loader_mapping:
            raise ValueError(f"不支持的文件格式: {ext}")

        # 使用对应的加载器加载文档
        loader = self.loader_mapping[ext](file_path)
        return loader.load()

    def load_documents_from_directory(self, directory_path: str, recursive: bool = True) -> list:
        """
        从目录加载多个文档
        :param directory_path: 目录路径
        :param recursive: 是否递归加载子目录
        :return: 所有文档列表
        """
        if not os.path.isdir(directory_path):
            raise NotADirectoryError(f"目录不存在: {directory_path}")

        documents = []
        for root, _, files in os.walk(directory_path):
            for file in files:
                file_path = os.path.join(root, file)
                try:
                    documents.extend(self.load_single_document(file_path))
                    print(f"成功加载文档: {file_path}")
                except Exception as e:
                    print(f"加载文档失败 {file_path}: {str(e)}")
            if not recursive:
                break

        return documents

if __name__ == "__main__":
    # 示例用法
    loader = DocumentLoader()
    try:
        # 加载单个PDF文档示例
        # single_doc = loader.load_single_document("example.pdf")
        # print(f"加载单个文档: {len(single_doc)}页")

        # 从目录加载所有文档示例
        docs = loader.load_documents_from_directory("../../examples/docs")
        print(f"从目录加载文档总数: {len(docs)}")
    except Exception as e:
        print(f"加载文档时出错: {str(e)}")